import os

import copy

import pandas as pd
import sql_metadata

from config import SQL_KEYWORDS, SQL_FUNCTIONS, KEYWORD_TYPE

## todo 该方法不完善需要进一步优化
def get_sql_skeleton(sql):
    '''
    Obtain the skeleton of SQL by removing non-keyword and non-operator tokens in SQL and replacing commas in SQL with XXX

    Attention：

    Args:
        sql: "SELECT sex, name FROM syjglt_tzzyczzxx"

    Returns:
        skeleton: String, "SELECT XXX FROM "

    '''

    tokens = sql_metadata.Parser(sql.upper()).tokens
    token_values = [i.value for i in tokens]

    remain_tokens = set(token_values) - set(SQL_KEYWORDS) - set(SQL_FUNCTIONS)

    skeleton = copy.deepcopy(token_values)

    for i in token_values:
        if i in remain_tokens and i != ',':
            skeleton.remove(i)
    skeleton = replace_comma_with_xxx(skeleton)
    skeleton = ' '.join(skeleton)

    return skeleton


def replace_comma_with_xxx(lst):
    for i in range(0, len(lst) - 1):
        if lst[i] == ',':
            lst[i] = 'XXX'

    return lst


def get_sql_keywords(sql):
    '''
    Get the keywords of SQL by parsing the SQL and removing non-keyword tokens in SQL
    Args:
        sql:String, "SELECT sex, FROM syjglt_tzzyczzxx"

    Returns:
        sql_keywords: Set,{'SELECT', 'FROM'}
    '''
    tokens = sql_metadata.Parser(sql.upper()).tokens
    token_values = [i.value for i in tokens]

    sql_keywords = set(token_values) & set(SQL_KEYWORDS) - set(SQL_FUNCTIONS)

    return sql_keywords


class SkeletonModel:
    skeleton_type_set = {'time_opt', 'and_or_where', 'except_set', 'window_fun', 'single_select', 'single_agg',
                         'concat_opt', 'exists_where', 'in_join', 'math_opt', 'single_group_by', 'multi_agg',
                         'not_where',
                         'distinct_select', 'in_where', 'union_set', 'arithmetic_opt', 'desc_order_by',
                         'multi_order_by',
                         'intersect_set', 'like_where', 'case', 'asc_order_by', 'having_group_by', 'single_order_by',
                         'table_sub_select', 'normal_where', 'multi_select', 'out_join', 'column_sub_select',
                         'limit_select', 'str_opt'}

    def __init__(self, name, sql):
        self.name = name  ## 骨架的名称
        self.total_count = 1  ## 骨架出现的总次数
        self.total_sql = [sql]  ## 包含该骨架的sql数
        self.len = len(name.split())  ## 骨架的长度
        self.keyword_fun_count = len(name.replace('XXX ', '').split()),  ## 骨架中包含的关键字、操作符和函数的个数

        self.type = get_skeleton_type(name)  ## 骨架类型

    def __str__(self):
        return f"""Name: {self.name}, Total Count: {self.total_count}, Total SQL: {self.total_sql}"""

    def update(self, sql):
        self.total_count += 1
        self.total_sql.append(sql)


class KeywordModel:

    def __init__(self, name):
        self.name = name  ## 关键字名称
        self.total_count = 1  ## 关键字出现的总次数
        self.total_sql = 1  ## 包含该关键字的sql数量
        self.max_times_in_sql = 1  ## 关键字在sql中出现的最大次数
        self.type = categorize_sql_keyword(name)  # 关键字的类型

    def __str__(self):
        return f"Name: {self.name}, Total Count: {self.total_count}, Total SQL: {self.total_sql}, Max Times in SQL: {self.max_times_in_sql}"

    def update(self, sql):
        ## update total_count and max_times_in_sql
        tokens = sql_metadata.Parser(sql.upper()).tokens
        self.total_sql += 1
        times_in_sql = 0
        for token in tokens:
            if token.value == self.name:
                times_in_sql += 1
                self.total_count += 1

        self.max_times_in_sql = times_in_sql if times_in_sql > self.max_times_in_sql else self.max_times_in_sql


def get_skeleton_type(name):
    '''
    Classify skeletons
    Args:
        name: String, 'SELECT DISTINCT FROM JOIN ON ='

    Returns:
        String, 'single_group_by'
    '''
    # print(name)
    skeleton_type = []
    name_list = name.split(' ')
    name_set = set(name_list)

    if 'WHERE' in name:

        if set('=, !=, <>, <, <=, !<, >, >=, !>, BETWEEN, IS NULL'.split(', ')) & name_set:
            skeleton_type.append('normal_where')
        if set(['IN', 'NOT IN']) & name_set:
            skeleton_type.append('in_where')
        if set(['LIKE', 'NOT LIKE']) & name_set:
            skeleton_type.append('like_where')
        if set(['NOT', 'IS']) & name_set and 'NOT IN' not in name_set and 'NOT LIKE' not in name_set:
            skeleton_type.append('not_where')
        if set(['AND', 'OR']) & name_set:
            skeleton_type.append('and_or_where')
        if set(['EXISTS']) & name_set:
            skeleton_type.append('exists_where')

    if 'ORDER BY' in name:
        if 'ORDER BY XXX' in name:
            skeleton_type.append('multi_order_by')
        elif 'ORDER BY XXX' not in name:
            skeleton_type.append('single_order_by')
        else:
            skeleton_type.append('other_order_by')
        if 'DESC' in name:
            skeleton_type.append('desc_order_by')
        elif 'DESC' not in name:
            skeleton_type.append('asc_order_by')

    if 'GROUP BY' in name:
        if 'HAVING' not in name_set:
            skeleton_type.append('single_group_by')
        if 'HAVING' in name_set:
            skeleton_type.append('having_group_by')
    if name_list.count('SELECT') >= 2 and 'UNION' not in name_set and 'JOIN' not in name_set:

        if 'FROM SELECT' in name:
            skeleton_type.append('table_sub_select')
        else:
            skeleton_type.append('column_sub_select')
    if 'JOIN' in name_set:
        if any(i in name for i in 'LEFT JOIN, LEFT OUTER JOIN, RIGHT JOIN, RIGHT OUTER JOIN'.split(',')):
            skeleton_type.append('out_join')
        elif any(i in name for i in 'JOIN, INNER JOIN'.split(',')):
            skeleton_type.append('in_join')
            # print(skeleton_type)
    if 'UNION' in name_set:
        skeleton_type.append('union_set')
    if 'INTERSECT' in name:
        skeleton_type.append('intersect_set')
    if 'EXCEPT' in name:
        skeleton_type.append('except_set')
    if any(i in name for i in 'OVER, RANK, ROW_NUMBER'.split(',')):
        skeleton_type.append('window_fun')
    if any(i in name for i in 'MIN, MAX, COUNT, SUM, AVG'.split(',')):
        if [i in name for i in 'MIN, MAX, COUNT, SUM, AVG'.split(',')].count(True) >= 2:
            skeleton_type.append('multi_agg')
        else:
            skeleton_type.append('single_agg')
    if set(['CONCAT', 'GROUP_CONCAT']) & name_set:
        skeleton_type.append('concat_opt')
    if any(i in name for i in ' +, -, *, /'.split(', ')):
        skeleton_type.append('arithmetic_opt')
    if any(i in name_set for i in 'UPPER, LOWER'.split(', ')):
        skeleton_type.append('str_opt')
    if any(i in name_set for i in 'NOW, YEAR, MONTH, EXTRACT, DATE'.split(', ')):
        skeleton_type.append('time_opt')
    if any(i in name_set for i in 'ABS, EXP, CEIL,FLOOR,ROUND,LOG'.split(', ')):
        skeleton_type.append('math_opt')
    if any(i in name_set for i in 'CASE, WHEN, THEN, ELSE, END'.split(', ')):
        skeleton_type.append('case')
    if len(skeleton_type) == 0 and "SELECT" in name:
        if 'XXX' in name:
            skeleton_type.append('multi_select')
        else:
            skeleton_type.append('single_select')
        if 'DISTINCT' in name:
            skeleton_type.append('distinct_select')
        if 'LIMIT' in name:
            skeleton_type.append('limit_select')
        if 'OFFSET' in name:
            skeleton_type.append('offset_select')
    return skeleton_type


def categorize_sql_keyword(keyword):
    '''
    Categorize keywords
    Args:
        keyword:  String, "WHERE"

    Returns:
        String, "筛选关键字"
    '''
    keyword = keyword.upper()
    if keyword in KEYWORD_TYPE.get('MATCH_KEYWORDS'):
        return "筛选关键字"
    elif keyword in ['SELECT', 'AS', 'FROM']:
        return keyword
    elif keyword in KEYWORD_TYPE.get('RANKING_KEYWORDS'):
        return "排序关键字"
    elif keyword in KEYWORD_TYPE.get('AGGREGATION_KEYWORDS'):
        keyword = keyword.upper()
    if keyword in KEYWORD_TYPE.get('MATCH_KEYWORDS'):
        return "筛选关键字"
    elif keyword in ['SELECT', 'AS', 'FROM']:
        return keyword
    elif keyword in KEYWORD_TYPE.get('RANKING_KEYWORDS'):
        return "排序关键字"
    elif keyword in KEYWORD_TYPE.get('AGGREGATION_KEYWORDS'):
        return "聚合关键字"
    elif keyword in KEYWORD_TYPE.get('JOIN_KEYWORDS'):
        return "关联关键字"
    elif keyword in KEYWORD_TYPE.get('SET_KEYWORD'):
        return "集合关键字"
    elif keyword in KEYWORD_TYPE.get('DATA_TYPE'):
        return "数据类型和格式转换关键字"
    elif keyword in KEYWORD_TYPE.get('BRANCH_KEYWORDS'):
        return "控制流关键字"
    elif keyword in KEYWORD_TYPE.get('OTHERS'):
        return "其他"
    else:
        return "Unknown"


def get_skeleton_models(sqls):
    '''
    Enter a set of sql and return the SkeletonModel information corresponding to the sql.
    Args:
        sqls: List[String], ['select distinct d.* from dept d join emp e on d.deptno = e.deptno;']

    Returns:
        skeleton_models: Dict{String: SkeletonModel} ,{'SELECT DISTINCT FROM JOIN ON =': SkeletonModel}
        skeleton_types: List[List[String]], [[ 'concat_opt', 'exists_where', 'in_join', 'math_opt'],[ 'concat_opt', 'exists_where']]
    '''
    skeleton_models = {}
    skeleton_types = []

    for sql in sqls:
        try:
            skeleton = get_sql_skeleton(sql)
        except:
            continue
        # print(skeleton)
        if not skeleton_models.get(skeleton):
            skeleton_model = SkeletonModel(skeleton, sql)
            skeleton_models[skeleton] = skeleton_model
        else:
            skeleton_model = skeleton_models[skeleton]
            skeleton_model.update(sql)
        # 获取每个sql的骨架类型
        skeleton_types.append(skeleton_model.type)
    return skeleton_models, skeleton_types


def get_keyword_models(sqls):
    '''
    Enter a set of sql and return the KeywordModel information corresponding to the sql.
    Args:
        sqls: List[String], ['select distinct d.* from dept d join emp e on d.deptno = e.deptno;']

    Returns:
        keyword_models: Dict[String, KeywordModel], {'distinct': KeywordModel1, 'join': KeywordModel2}
    '''
    keyword_models = {}

    for sql in sqls:
        try:
            keywords = get_sql_keywords(sql)
        except:
            continue
        for keyword in keywords:
            if not keyword_models.get(keyword):
                keyword_model = KeywordModel(name=keyword)
                keyword_models[keyword] = keyword_model
            else:
                keyword_model = keyword_models.get(keyword)
                keyword_model.update(sql)
    return keyword_models


def save_skeleton_models(skeleton_models, file_name):
    '''
    Enter the SkeletonModel of the skeleton and the file storage address to store the keyword statistics.
    Args:
        skeleton_models: Dict{String: SkeletonModel} ,{'SELECT DISTINCT FROM JOIN ON =': SkeletonModel}
        file_name: String, "langchao_skeleton_data.xlsx"

    Returns:
        None
    '''
    skeleton_data = [{
        '骨架格式': model.name,
        '骨架总个数': model.total_count,
        '使用该骨架的sql': model.total_sql,
        'example': model.total_sql[0],
        'type': model.type,
        '骨架长度': model.len,
        '骨架包含的关键字、方法、操作符个数': len(model.name.replace('XXX ', '').split()),
        '骨架包含的关键字方法和操作符': model.name.replace('XXX ', '').split(),
    } for model in skeleton_models.values()]
    data_frame = pd.DataFrame(skeleton_data)
    data_frame.to_excel(file_name, index=False)
    print('skeleton_models saved!')


def save_keyword_models(keyword_models, file_name):
    '''
    Enter the keyword's KeywordModel and the file storage address to store the keyword's statistical information.
    Args:
        keyword_models: Dict[String, KeywordModel], {'distinct': KeywordModel1, 'join': KeywordModel2}
        file_name: String, "langchao_keyword_data.xlsx"

    Returns:
        None
    '''
    keyword_data = [{
        '关键字类型': model.type,
        '关键字名称': model.name,
        '关键字出现的总次数': model.total_count,
        '包含该关键字的sql数量': model.total_sql,
        '关键字在一个sql中出现的最多次数': model.max_times_in_sql,
    } for model in keyword_models.values()]
    data_frame = pd.DataFrame(keyword_data)

    data_frame.to_excel(file_name, index=False)
    print('keyword_models saved!')


def save_skeleton_type(skeleton_types, file_name):
    '''
    Enter a list of skeleton types and the path where the file needs to be stored to store the skeleton type.

    Args:
        skeleton_types: List[List[String]], [[ 'concat_opt', 'exists_where', 'in_join', 'math_opt'],[ 'concat_opt', 'exists_where']]
        file_name: String, "langchao_skeleton_type_data.xlsx"

    Returns:
        None
    '''

    def get_chinese_skeleton_type(keyword):
        ## 基本检索：xxx_select; 排序：xxx_order_by; 筛选：xxx_where
        ## 分组： xxx_group_by; 子查询：sub_select; 多表联结：xxx_join
        ## 集合操作： xxx_set; 窗口函数：window_fun; 聚合：xxx_agg
        ## 计算： xxx_opt; 控制流函数：case
        if keyword == 'normal_where':
            return '筛选', '基本筛选', 'WHERE, =, !=, <>, <, <=, !<, >, >=, !>, BETWEEN, IS NULL', 'SELECT prod_name, prod_price FROM Products WHERE prod_price <= 10;'
        elif keyword == 'in_where':
            return '筛选', 'IN 筛选', 'WHERE, IN', "SELECT prod_name, prod_price FROM Products WHERE vend_id IN ('DLL01','BRS01');"
        elif keyword == 'like_where':
            return '筛选', 'LIKE 筛选', 'WHERE，LIKE', "SELECT prod_id, prod_name FROM Products WHERE prod_name LIKE 'Fish%';"
        elif keyword == 'not_where':
            return '筛选', 'NOT,IS 筛选', 'WHERE, NOT, IS', "SELECT prod_name FROM Products WHERE NOT vend_id = 'DLL01';"
        elif keyword == 'and_or_where':
            return '筛选', 'AND,OR 筛选', 'WHERE，AND，OR', "SELECT prod_id, prod_price, prod_name FROM Products WHERE vend_id = 'DLL01' AND prod_price <= 4;"
        elif keyword == 'exists_where':
            return '筛选', 'EXISTS 筛选', 'WHERE, EXISTS', "SELECT prod_name FROM Products WHERE EXISTS (SELECT prod_name FROM Products WHERE vend_id = dll01);"
        elif keyword == 'multi_order_by':
            return '排序', '多字段排序', 'ORDER BY', "SELECT prod_id, prod_name FROM Products ORDER BY prod_price, prod_name;"
        elif keyword == 'single_order_by':
            return '排序', '单字段排序', 'ORDER BY', "SELECT prod_id, prod_name FROM Products ORDER BY prod_name;"
        elif keyword == 'desc_order_by' or keyword == 'asc_order_by':
            return '排序', '降序排序或升序排列', 'ORDER BY,DESC,ASC(默认)', "SELECT prod_id, prod_name FROM Products ORDER BY prod_name DESC;"
        elif keyword == 'single_group_by':
            return '分组', '基本分组', 'GROUP BY', "SELECT prod_id, prod_name FROM Products GROUP BY prod_name;"
        elif keyword == 'having_group_by':
            return '分组', '过滤分组', 'GROUP BY', "SELECT cust_id, COUNT(＊) AS orders FROM Orders GROUP BY cust_id HAVING COUNT(＊) >= 2;"
        elif keyword == 'table_sub_select':
            return '子查询', '用作表子查询', 'as table', "SELECT avg(price) FROM (SELECT price FROM products WHERE category = Electronics) AS electronics;"
        elif keyword == 'column_sub_select':
            return '子查询', '用作列子查询', 'as column', "SELECT prod_name FROM Products WHERE prod_id = (SELECT prod_id FROM Products WHERE prod_name = 'Chai');"
        elif keyword == 'out_join':
            return '联结', '外联结', 'LEFT JOIN, LEFT OUTER JOIN, RIGHT JOIN, RIGHT OUTER JOIN', "SELECT Customers.cust_id, Orders.order_num FROM Customers LEFT OUTER JOIN Orders ON Customers.cust_id = Orders.cust_id;"
        elif keyword == 'in_join':
            return '联结', '内联结', 'JOIN, INNER JOIN', "SELECT vend_name, prod_name, prod_price FROM Vendors INNER JOIN Products ON Vendors.vend_id = Products.vend_id;"
        elif keyword == 'union_set':
            return '组合查询', 'UNION 组合查询', 'UNION', " SELECT cust_name, cust_contact, cust_email FROM Customers WHERE cust_state IN ('IL','IN','MI') UNION SELECT cust_name, cust_contact, cust_email FROM Customers WHERE cust_name = 'Fun4All';"
        elif keyword == 'intersect_set':
            return '组合查询', 'INTERSECT 组合查询', 'INTERSECT', "SELECT a.accident_id, a.accident_date, a.accident_location, v.vehicle_id, v.license_plate, v.owner_name FROM accidents a JOIN vehicles v ON a.vehicle_id = v.vehicle_id JOIN highways h ON v.highway_id = h.highway_id WHERE a.accident_location = 'Intersection B' AND h.highway_name = 'Highway B';"
        elif keyword == 'except_set':
            return '组合查询', 'EXCEPT 组合查询', 'except', "SELECT * FROM Customers WHERE cust_name = 'Fun4All' EXCEPT SELECT * FROM Customers WHERE cust_name = 'Fun4"
        elif keyword == 'window_fun':
            return '窗口函数', '窗口函数', 'OVER, RANK, ROW_NUMBER等', "SELECT fiscal_year, sales_employee, sale, SUM(sale) OVER (PARTITION BY fiscal_year) total_sales FROM sales; "
        elif keyword == 'case':
            return '控制流关键字', 'CASE 关键字', 'CASE, WHEN, THEN, ELSE, END等', "SELECT c.psn_name as 姓名, c.certno as 身份证号, c.psn_no as 个人编号, (case when max(a.psn_type) like '11%' then '在职职工' when max(a.psn_type) like '12%' then '退休职工' else '' end) as 参保身份 FROM staf_psn_clct_detl_d a INNER JOIN psn_info_b c ON a.psn_no = c.psn_no WHERE a.accrym_begn >= '201911' AND a.clct_flag = '1' AND a.revs_flag = 'Z' group by c.psn_name"
        elif keyword == 'multi_agg':
            return '聚合函数', '多字段聚合函数', '多个基本聚合函数', "SELECT MIN(prod_price) AS min_price, MAX(prod_price) AS max_price, SUM(prod_price) from  Products;"
        elif keyword == 'single_agg':
            return '聚合函数', '单字段聚合函数', 'MIN, MAX, COUNT, SUM, AVG', "SELECT AVG(prod_price) AS avg_price FROM Products;"
        elif keyword == 'concat_opt':
            return '计算', '字符串连接函数', 'Concat, Group Concat', "SELECT Concat(vend_name, ' (', vend_country, ')') FROM Vendors ORDER BY vend_name;"
        elif keyword == 'arithmetic_opt':
            return '计算', '算术运算符', ' +, -, *, /', "SELECT prod_id, quantity, item_price, quantity＊item_price AS expanded_price FROM OrderItems WHERE order_num = 20008;"
        elif keyword == 'str_opt':
            return '计算', '字符串运算符', 'UPPER, LOWER等', "SELECT prod_id, prod_name, UPPER(prod_name) AS up_prod_name FROM Products;"
        elif keyword == 'time_opt':
            return '计算', '时间运算符', "NOW, YEAR, MONTH, EXTRACT, DATE等。完整见： https://www.runoob.com/mysql/mysql-functions.html ", "SELECT order_num FROM Orders WHERE YEAR(order_date) = 2020;"
        elif keyword == 'math_opt':
            return '计算', '数学运算符', 'ABS, EXP, CEIL,FLOOR,ROUND,LOG等', "SELECT prod_id, quantity, FLOOR(item_price) AS expanded_price FROM OrderItems WHERE order_num = 2"
        elif keyword == 'multi_select':
            return '基本检索', '多字段选择', '仅Select', "SELECT prod_id, prod_name, prod_price FROM Products;"
        elif keyword == 'single_select':
            return '基本检索', '单字段选择', '仅Select', "SELECT prod_name FROM Products;"
        elif keyword == 'distinct_select':
            return '基本检索', '去重选择', 'DISTINCT', "SELECT DISTINCT vend_id FROM Products;"
        elif keyword == 'limit_select' or keyword == 'offset_select':
            return '基本检索', '限制选择', 'LIMIT, OFFSET', "SELECT prod_name FROM Products LIMIT 5;"
        else:
            return keyword, keyword, keyword, keyword

    skeleton_type_dict = {}
    skeleton_type_set = SkeletonModel.skeleton_type_set
    for k in skeleton_type_set:
        count = 0
        for keyword_list in skeleton_types:
            count += keyword_list.count(k)
        skeleton_type_dict[k] = count

    keyword_data = []
    for k, v in skeleton_type_dict.items():
        first_type, second_type, keyword, example = get_chinese_skeleton_type(k)
        keyword_data.append({
            "一级类别": first_type,
            "二级类别": second_type,
            "关键字": keyword,
            "出现的次数": v,
            "例子": example,
        })

    data_frame = pd.DataFrame(keyword_data)

    data_frame.to_excel(file_name, index=False)
    print('keyword_models saved!')
